import pymysql
import pymysql.cursors
import time
import threading
from .config import get_database_config

class DatabaseConnector:
    """数据库连接类，负责处理数据库连接和查询操作"""
    
    def __init__(self):
        """初始化数据库连接器"""
        self.db_config = get_database_config()
        self.connection = None
        self.cursor = None
        self.heartbeat_interval = 30  # 心跳间隔，单位秒
        self.heartbeat_thread = None
        self.heartbeat_running = False
    
    def check_database_structure(self):
        """检查数据库结构是否存在"""
        # 检查所有必要的表是否存在
        tables = ['owners', 'owner_relations', 'ips', 'domains', 'ipdomains', 'assets']
        missing_tables = []
        
        try:
            for table in tables:
                query = f"SHOW TABLES LIKE '{table}'"
                self.cursor.execute(query)
                if not self.cursor.fetchone():
                    missing_tables.append(table)
            
            if missing_tables:
                print(f"警告: 缺少必要的数据库表: {', '.join(missing_tables)}")
                print("请先运行 mysql_db_generator.py 创建数据库结构")
        except Exception as e:
            print(f"检查数据库结构时出错: {str(e)}")
    
    def connect(self):
        """连接到数据库"""
        try:
            # 获取数据库配置
            self.db_config = get_database_config()
            
            # 连接数据库，添加更多参数保持连接活跃
            self.connection = pymysql.connect(
                host=self.db_config.get('host', 'localhost'),
                user=self.db_config.get('user', 'root'),
                password=self.db_config.get('password', ''),
                database=self.db_config.get('db_name', 'asset_collection'),
                port=int(self.db_config.get('port', 3306)),
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor,
                connect_timeout=10,
                read_timeout=30,
                write_timeout=30,
                autocommit=True,
                # 添加保持连接的参数
                client_flag=pymysql.constants.CLIENT.MULTI_RESULTS
            )
            self.cursor = self.connection.cursor()
            
            # 检查数据库和表是否存在
            self.check_database_structure()
            
            # 启动心跳线程
            self._start_heartbeat()
        except pymysql.MySQLError as e:
            error_code, error_message = e.args
            if error_code == 1045:
                print(f"数据库连接错误: 用户名或密码错误\n详细信息: {error_message}")
            elif error_code == 1049:
                print(f"数据库连接错误: 数据库 '{self.db_config.get('db_name', 'asset_collection')}' 不存在\n详细信息: {error_message}")
            elif error_code == 2003:
                print(f"数据库连接错误: 无法连接到MySQL服务器 '{self.db_config.get('host', 'localhost')}:{self.db_config.get('port', 3306)}'\n详细信息: {error_message}")
            else:
                print(f"数据库连接错误: {error_message}")
            self.connection = None
            self.cursor = None
        except ValueError as e:
            print(f"配置错误: {str(e)}")
            self.connection = None
            self.cursor = None
        except Exception as e:
            print(f"未知连接错误: {str(e)}")
            self.connection = None
            self.cursor = None
    
    def disconnect(self):
        """断开数据库连接"""
        # 停止心跳线程
        self._stop_heartbeat()
        
        if self.connection:
            try:
                if self.cursor:
                    self.cursor.close()
                self.connection.close()
            except:
                pass  # 忽略关闭时的错误
            finally:
                self.connection = None
                self.cursor = None
    
    def _start_heartbeat(self):
        """启动心跳线程，定期发送ping保持连接活跃"""
        if self.heartbeat_running:
            return
        
        self.heartbeat_running = True
        self.heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
        self.heartbeat_thread.start()
        
    def _stop_heartbeat(self):
        """停止心跳线程"""
        self.heartbeat_running = False
        if self.heartbeat_thread and self.heartbeat_thread.is_alive():
            self.heartbeat_thread.join(timeout=1.0)
        
    def _heartbeat_loop(self):
        """心跳循环，定期发送ping命令"""
        while self.heartbeat_running:
            try:
                time.sleep(self.heartbeat_interval)
                if self.connection and self.connection.open:
                    # 发送一个简单的ping查询
                    with self.connection.cursor() as temp_cursor:
                        temp_cursor.execute("SELECT 1")
            except Exception:
                # 心跳失败，不做处理，让执行查询时的重试机制处理
                pass
        
    def execute_query(self, query, params=None):
        """执行SQL查询，包含重试机制"""
        max_retries = 3
        retry_count = 0
        result = None
        
        while retry_count < max_retries:
            try:
                # 确保连接是活跃的
                if not self.connection or not self.connection.open:
                    self.connect()
                    if not self.connection:
                        return None
                
                # 执行查询前先检查连接是否真的可用
                try:
                    with self.connection.cursor() as temp_cursor:
                        temp_cursor.execute("SELECT 1")
                except:
                    # 连接不可用，重新连接
                    self.disconnect()
                    self.connect()
                    if not self.connection:
                        retry_count += 1
                        time.sleep(1)
                        continue
                
                # 执行查询
                self.cursor.execute(query, params or ())
                # 对于查询类型的语句，不需要commit
                if not query.strip().upper().startswith('SELECT'):
                    self.connection.commit()
                result = self.cursor
                break  # 查询成功，跳出循环
            except pymysql.MySQLError as e:
                error_code, error_message = e.args
                print(f"数据库操作错误: {str(e)}")
                
                # 只在连接仍然有效的情况下尝试rollback
                if self.connection and self.connection.open:
                    try:
                        self.connection.rollback()
                    except:
                        pass  # 如果rollback也失败，忽略错误
                
                # 针对连接断开的错误进行特殊处理
                if error_code in (2006, 2013, 2055, 2014):  # MySQL server has gone away, Lost connection, Lost connection to MySQL server at '%s', system error: %d, Commands out of sync
                    retry_count += 1
                    print(f"连接断开，尝试重新连接 ({retry_count}/{max_retries})...")
                    time.sleep(1)  # 等待1秒后重试
                    # 强制重新连接
                    self.disconnect()
                    try:
                        self.connect()
                    except Exception as re_conn_error:
                        print(f"重新连接失败: {re_conn_error}")
                else:
                    break  # 其他错误不重试
            except Exception as e:
                print(f"未知错误: {str(e)}")
                # 只在连接仍然有效的情况下尝试rollback
                if self.connection and self.connection.open:
                    try:
                        self.connection.rollback()
                    except:
                        pass
                break  # 未知错误不重试
        
        return result

    def fetchvalue(self,sql, params=None):
        result = self.execute_query(sql, params)
        return list(result.fetchone().values())[0] if result else None

    def fetchall(self, sql, params=None):
        result = self.execute_query(sql, params)
        return result.fetchall() if result else []

    def fetchone(self, sql, params=None):
        result = self.execute_query(sql, params)
        return result.fetchone() if result else None